package main

import (
	"context"
	"fmt"
	"log"
	"time"

	clientv3 "go.etcd.io/etcd/client/v3"
)

type DistributedQueue struct {
	client   *clientv3.Client
	queueKey string
}

func NewDistributedQueue(endpoints []string, queueKey string) (*DistributedQueue, error) {
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   endpoints,
		DialTimeout: 5 * time.Second,
	})
	if err != nil {
		return nil, err
	}

	return &DistributedQueue{
		client:   cli,
		queueKey: queueKey,
	}, nil
}

func (dq *DistributedQueue) Enqueue(item string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// Use a unique key for each item in the queue, e.g., appending a timestamp or UUID
	key := fmt.Sprintf("%s/%d", dq.queueKey, time.Now().UnixNano())
	_, err := dq.client.Put(ctx, key, item)
	return err
}

func (dq *DistributedQueue) Dequeue() (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// Get the first item in the queue (lexicographically sorted by key)
	resp, err := dq.client.Get(ctx, dq.queueKey, clientv3.WithPrefix(), clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend))
	if err != nil {
		return "", err
	}

	if len(resp.Kvs) == 0 {
		return "", fmt.Errorf("queue is empty")
	}

	// The first key-value pair in the sorted response is the oldest item in the queue
	oldestKey := string(resp.Kvs[0].Key)
	item := string(resp.Kvs[0].Value)

	// Delete the oldest item from the queue
	_, err = dq.client.Delete(ctx, oldestKey)
	if err != nil {
		return "", err
	}

	return item, nil
}

func (dq *DistributedQueue) Close() {
	dq.client.Close()
}

func main() {
	endpoints := []string{"localhost:2379"}
	queueKey := "/my-distributed-queue"

	queue, err := NewDistributedQueue(endpoints, queueKey)
	if err != nil {
		log.Fatalf("failed to create distributed queue: %v", err)
	}
	defer queue.Close()

	// Enqueue items
	err = queue.Enqueue("item1")
	if err != nil {
		log.Fatalf("failed to enqueue item: %v", err)
	}
	err = queue.Enqueue("item2")
	if err != nil {
		log.Fatalf("failed to enqueue item: %v", err)
	}

	// Dequeue items
	item, err := queue.Dequeue()
	if err != nil {
		log.Fatalf("failed to dequeue item: %v", err)
	}
	fmt.Printf("Dequeued: %s\n", item)

	item, err = queue.Dequeue()
	if err != nil {
		log.Fatalf("failed to dequeue item: %v", err)
	}
	fmt.Printf("Dequeued: %s\n", item)
}
